import vllm
from vllm import LLM, SamplingParams

print("=====generate=====")
path = "/data/mtt/model_convert/llama-2-70b-chat-hf-fp16-convert-tp4-new/"
# export MTT_LIB_PATH=/data/mtt/MT-Transformer/build/src/mttransformer/th_op/libmt_transformer_pt.so
llm = LLM(model=path, gpu_memory_utilization = 0.9, tensor_parallel_size = 4, device = "musa", block_size = 64, max_num_seqs = 128)
print("llm construct end")

# Print the outputs.
prompts = [
    "[INST] Hello! [/INST]",
    "[INST] Hi! who are you? [/INST]",
    "[INST] I am going to Paris, where should I go? [/INST]",
    "[INST] What is your name? [/INST]",
    "[INST] Where are you coming from? [/INST]",
  ]*30
# prompts.extend("[INST] Nice to meet you! [/INST]" for i in range(20))
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=64)
outputs = llm.generate(prompts, sampling_params)
print(f"\n")
for idx, output in enumerate(outputs):
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt {idx}: {prompt}, \nGenerated text: {generated_text}\n\n")
